{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "import sys\n",
    "import numpy as np\n",
    "import argparse\n",
    "from scipy.spatial.transform import Rotation as R \n",
    "import open3d as o3d \n",
    "# add catkin_ws context \n",
    "sys.path.append(\"/home/junting/catkin_ws/devel/lib/python3.9/site-packages\")\n",
    "sys.path.append(\"/home/junting/franka_ws/devel/lib/python3.9/site-packages\")\n",
    "\n",
    "from src.lmp import *\n",
    "from src.env.true_grounding_env import TrueGroundingEnv\n",
    "from src.configs.config import load_config\n",
    "cfg_tabletop = load_config(\"perception_few_shot_gpt_3.5.yaml\")\n",
    "import rospy \n",
    "import rospkg\n",
    "import jupyros as jr\n",
    "\n",
    "from std_msgs.msg import String, Header\n",
    "from geometry_msgs.msg import PoseStamped, Pose, Point, Quaternion\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "########################################\n",
    "# initialize environment\n",
    "########################################\n",
    "rospy.init_node('eval_code', log_level=rospy.DEBUG)\n",
    "# get package root path \n",
    "pkg_root = rospkg.RosPack().get_path('instruct_to_policy')\n",
    "\n",
    "# setup environment\n",
    "env = TrueGroundingEnv(cfg_tabletop)\n",
    "# env.reset()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ground truth perception and 3D fusion "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3D fusion \n",
    "from src.perception.scene_manager import SceneManager\n",
    "sensor_data = env.get_sensor_data()\n",
    "sensor_data['detections_list'] = [{},{},{}]\n",
    "\n",
    "scene_manager = SceneManager()\n",
    "scene_manager.update_fusion(sensor_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get cropped point cloud\n",
    "# drawer_bbox = env.get_3d_bbox('cabinet.drawer0')\n",
    "# drawer_pcd = scene_manager.scene_tsdf_full.crop_cloud(\n",
    "#     crop_center=(drawer_bbox[:3] + drawer_bbox[3:]) / 2,\n",
    "#     crop_size=(drawer_bbox[3:] - drawer_bbox[:3])\n",
    "# )\n",
    "# cabinet_bbox = env.get_3d_bbox('cabinet')\n",
    "# cabinet_pcd = scene_manager.scene_tsdf_full.crop_cloud(\n",
    "#     crop_center=(cabinet_bbox[:3] + cabinet_bbox[3:]) / 2,\n",
    "#     crop_size=(cabinet_bbox[3:] - cabinet_bbox[:3])\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import open3d as o3d \n",
    "# o3d.io.write_point_cloud(os.path.join(pkg_root, 'data', 'drawer.pcd'), drawer_pcd)\n",
    "# o3d.io.write_point_cloud(os.path.join(pkg_root, 'data', 'cabinet.pcd'), cabinet_pcd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pick and Place"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "pose_msg = env.parse_adaptive_shape_grasp_pose(\"apple\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rot_mat = R.from_quat([pose_msg.orientation.x, pose_msg.orientation.y, pose_msg.orientation.z, pose_msg.orientation.w]).as_matrix()\n",
    "translation = np.array([pose_msg.position.x, pose_msg.position.y, pose_msg.position.z])\n",
    "\n",
    "depth = 0.05\n",
    "\n",
    "pregrasp_offset_local = np.array([0, 0, -0.15]).astype(np.float32)\n",
    "# predicted gripper center is 0.02m above the gripper tip\n",
    "approach_offset_local = np.array([0, 0, depth - 0.02 ]).astype(np.float32)\n",
    "pregrasp_position = translation + rot_mat @ pregrasp_offset_local\n",
    "approach_position = translation + rot_mat @ approach_offset_local\n",
    "\n",
    "pregrasp_pose = Pose(position=Point(*pregrasp_position), orientation=pose_msg.orientation)\n",
    "approach_pose = Pose(position=Point(*approach_position), orientation=pose_msg.orientation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.open_gripper()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.publish_goal_to_marker(pregrasp_pose)\n",
    "env.move_to_pose(pregrasp_pose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.publish_goal_to_marker(approach_pose)\n",
    "env.move_to_pose(approach_pose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.close_gripper(width=0.05, force=30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.attach_object(\"apple\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "place_pose = env.parse_place_pose(object_name=\"apple\", receptacle_name=\"white_ceramic_plate\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.move_to_pose(place_pose)\n",
    "env.open_gripper()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.detach_object(\"apple\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Open Drawer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jr.publish('/rviz/moveit/move_marker/goal_panda_hand_tcp', PoseStamped)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.reset()\n",
    "# env.open_gripper()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# env.get_obj_name_list()\n",
    "# [bbox.object_id for bbox in env.gazebo_gt_bboxes]\n",
    "env.get_3d_bbox(\"cabinet.handle_0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grasp_pose = env.parse_adaptive_shape_grasp_pose(object_name=\"cabinet.handle_0\")\n",
    "print(grasp_pose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.publish_goal_to_marker(grasp_pose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.grasp(grasp_pose)\n",
    "# env.move_to_pose(grasp_pose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.close_gripper()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate a horizontal trajectory to open the drawer\n",
    "grasp_position = np.array([grasp_pose.position.x, grasp_pose.position.y, grasp_pose.position.z])\n",
    "pull_position = grasp_position + np.array([0.2, 0, 0]).astype(float)\n",
    "pull_pose = Pose(position=Point(*pull_position), orientation=grasp_pose.orientation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.move_to_pose(pull_pose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.open_gripper()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test different grasp depth \n",
    "pose_msg = env.parse_adaptive_shape_grasp_pose(object_name=\"apple\")\n",
    "\n",
    "rot_mat = R.from_quat([pose_msg.orientation.x, pose_msg.orientation.y, pose_msg.orientation.z, pose_msg.orientation.w]).as_matrix()\n",
    "translation = np.array([pose_msg.position.x, pose_msg.position.y, pose_msg.position.z])\n",
    "\n",
    "depth = 0.05\n",
    "\n",
    "pregrasp_offset_local = np.array([0, 0, -0.15]).astype(np.float32)\n",
    "# predicted gripper center is 0.02m above the gripper tip\n",
    "approach_offset_local = np.array([0, 0, depth - 0.02 ]).astype(np.float32)\n",
    "pregrasp_position = translation + rot_mat @ pregrasp_offset_local\n",
    "approach_position = translation + rot_mat @ approach_offset_local\n",
    "\n",
    "pregrasp_pose = Pose(position=Point(*pregrasp_position), orientation=pose_msg.orientation)\n",
    "approach_pose = Pose(position=Point(*approach_position), orientation=pose_msg.orientation)\n",
    " \n",
    "env.open_gripper()\n",
    " \n",
    "# env.publish_goal_to_marker(pregrasp_pose)\n",
    "env.move_to_pose(pregrasp_pose)\n",
    "# env.publish_goal_to_marker(approach_pose)\n",
    "env.move_to_pose(approach_pose)\n",
    "env.close_gripper(width=0.05, force=30)\n",
    "env.attach_object(\"apple\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.reset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Grasp preference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from graspnetAPI import GraspGroup, Grasp\n",
    "from grasp_detection.msg import Grasp as GraspMsg\n",
    "from src.grasp_detection.utils import select_grasp_by_preference\n",
    "\n",
    "# reload src.grasp_detection.utils \n",
    "import importlib\n",
    "importlib.reload(sys.modules['src.grasp_detection.utils'])\n",
    "from src.grasp_detection.utils import select_grasp_by_preference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "object_name = \"white_and_pink_box\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ground truth preference of object \n",
    "\n",
    "# preferred_position is the center of the object\n",
    "preferred_position = env.get_object_center_position(object_name)\n",
    "\n",
    "# preferred_orientation is table surface \n",
    "preferred_plane_normal = np.array([0, 0, 1])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Manually call grasp detection model to get list of grasps\n",
    "\n",
    "object_bbox = env.get_3d_bbox(object_name)    \n",
    "sensor_data = env.get_sensor_data()\n",
    "\n",
    "object_pcd = scene_manager.scene_tsdf_full.crop_cloud(\n",
    "    crop_center=(object_bbox[:3] + object_bbox[3:]) / 2,\n",
    "    crop_size=(object_bbox[3:] - object_bbox[:3])\n",
    ")\n",
    "\n",
    "bbox_center = (object_bbox[:3] + object_bbox[3:]) / 2\n",
    "bbox_size = object_bbox[3:] - object_bbox[:3]\n",
    "        \n",
    "data = {\n",
    "    # 'detections_list': detections_list,\n",
    "    'bboxes_3d_dict':{\n",
    "        object_name:{'center': bbox_center, 'size': bbox_size}\n",
    "    }\n",
    "}\n",
    "data.update(sensor_data)\n",
    "\n",
    "# call grasp detection service\n",
    "grasp_candidates: List[GraspMsg] = env.grasp_model.predict(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Manually select a grasp from the list of grasps\n",
    "grasp_idx, weighted_score = select_grasp_by_preference(grasp_candidates, preferred_position=preferred_position, preferred_plane_normal=preferred_plane_normal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize all grasp candidates with object point cloud \n",
    "grasp_o3d_meshes = []\n",
    "\n",
    "# rotate grasp orientation from gazebo gripper frame to anygrasp gripper frame\n",
    "rot_anygrasp2gazebo = np.array([[0,0,1],[0,1,0],[-1,0,0]])\n",
    "rot_gazebo2anygrasp = np.linalg.inv(rot_anygrasp2gazebo)\n",
    "\n",
    "for i, grasp_msg in enumerate(grasp_candidates):\n",
    "    quat = np.array([grasp_msg.grasp_pose.orientation.x, grasp_msg.grasp_pose.orientation.y, grasp_msg.grasp_pose.orientation.z, grasp_msg.grasp_pose.orientation.w])\n",
    "    rotation_matrix = R.from_quat(quat).as_matrix()\n",
    "    rotation_matrix = rotation_matrix @ rot_gazebo2anygrasp\n",
    "    translation = np.array([grasp_msg.grasp_pose.position.x, grasp_msg.grasp_pose.position.y, grasp_msg.grasp_pose.position.z])\n",
    "    # [score, width, height, depth, rotation_matrix, translation, object_id]\n",
    "    grasp = Grasp(\n",
    "        *[grasp_msg.grasp_score, grasp_msg.grasp_width, 0.02, grasp_msg.grasp_depth, rotation_matrix, translation, i]\n",
    "    )\n",
    "    grasp_o3d_meshes.append(grasp.to_open3d_geometry())\n",
    "    \n",
    "# change the selected grasp color to red \n",
    "grasp_o3d_meshes[grasp_idx].paint_uniform_color([1,0,0])\n",
    "\n",
    "o3d.visualization.draw_geometries([object_pcd, *grasp_o3d_meshes])    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# predict the grasp pose of white_and_pink_box with anygrasp\n",
    "pose_msg = env.parse_adaptive_shape_grasp_pose(object_name=object_name, preferred_position=preferred_position, preferred_plane_normal=preferred_plane_normal)\n",
    "env.publish_goal_to_marker(pose_msg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Joint Prediction "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.get_gt_bbox('cabinet.handle_0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get handle 0 position \n",
    "handle_0_position = env.get_object_center_position(\"cabinet.handle_0\")\n",
    "print(handle_0_position)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "joint = env.get_object_joint_info(obj_name=\"cabinet\", position=handle_0_position, type=\"any\")\n",
    "joint_axis = joint['joint_axis']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.open_gripper(width=0.10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# preferred_position: Optional(np.ndarray), prefered gripper tip point position \n",
    "# preferred_approach_direction: Optional(np.ndarray), prefered gripper approach direction\n",
    "# preferred_plane_normal: Optional(np.ndarray), prefered gripper plane normal direction  \n",
    "\n",
    "grasp_pose = env.parse_adaptive_shape_grasp_pose(object_name=\"cabinet.handle_0\", preferred_position=handle_0_position, preferred_approach_direction=joint_axis)\n",
    "env.publish_goal_to_marker(grasp_pose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.grasp(grasp_pose, pre_grasp_approach=0.05, tentative_depth_list=[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.close_gripper()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate a horizontal trajectory to open the drawer\n",
    "def move_in_direction(axis: np.array, distance: float):\n",
    "    current_pose = env.get_gripper_pose()\n",
    "    target_pose = Pose()\n",
    "    normalized_axis = np.array(axis) / np.linalg.norm(axis)\n",
    "    target_pose.position.x = axis[0] * distance + current_pose.position.x\n",
    "    target_pose.position.y = axis[1] * distance + current_pose.position.y\n",
    "    target_pose.position.z = axis[2] * distance + current_pose.position.z\n",
    "    target_pose.orientation = current_pose.orientation\n",
    "    env.move_to_pose(target_pose)\n",
    "    \n",
    "move_in_direction(axis=joint_axis, distance=0.2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ROS Python 3 (ipykernel)",
   "language": "python",
   "name": "ros_python"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
